import random
import sys
import json
import os
from datetime import datetime
from openai import OpenAI
from dotenv import load_dotenv, dotenv_values

# 设置控制台输出编码为 UTF-8
sys.stdout.reconfigure(encoding="utf-8")


class WeatherTimeTools:
    @staticmethod
    def get_current_weather(arguments):

        weather_conditions = ["晴天", "多云", "雨天"]
        random_weather = random.choice(weather_conditions)
        location = arguments["location"]
        return f"{location}今天是{random_weather}。"

    @staticmethod
    def get_current_time():
        current_datetime = datetime.now()
        formatted_time = current_datetime.strftime("%Y-%m-%d %H:%M:%S")
        return f"当前时间：{formatted_time}。"


class AssistantConfig:
    TOOLS = [
        {
            "type": "function",
            "function": {
                "name": "get_current_time",
                "description": "当你想知道现在的时间时非常有用。",
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "当你想查询指定城市的天气时非常有用。",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "城市或县区，比如北京市、杭州市、余杭区等。",
                        }
                    },
                    "required": ["location"],
                },
            },
        },
    ]

    MESSAGES = [
        {
            "role": "system",
            "content": """你是一个很有帮助的助手。如果用户提问关于天气的问题，请调用 'get_current_weather' 函数;
            如果用户提问关于时间的问题，请调用'get_current_time'函数。
            请以友好的语气回答问题。""",
        },
        {"role": "user", "content": "上海天气"},
    ]


class Assistant:

    # 定义环境变量名称字符串
    DASHSCOPE_API_KEY = "DASHSCOPE_API_KEY"
    DASHSCOPE_BASE_URL = "DASHSCOPE_BASE_URL"

    def __init__(self):
        # 加载.env文件
        load_dotenv()
        env_vars = dotenv_values()

        # 检查环境变量是否存在
        if not env_vars.get(self.DASHSCOPE_API_KEY):
            raise ValueError("请在.env文件中设置 DASHSCOPE_API_KEY")

        self.client = OpenAI(
            api_key=env_vars.get(self.DASHSCOPE_API_KEY),
            base_url=env_vars.get(self.DASHSCOPE_BASE_URL),
        )
        self.tools = WeatherTimeTools()
        self.function_mapper = {
            "get_current_weather": self.tools.get_current_weather,
            "get_current_time": self.tools.get_current_time,
        }

    def function_calling(self):
        completion = self.client.chat.completions.create(
            model="qwen-plus",
            messages=AssistantConfig.MESSAGES,
            tools=AssistantConfig.TOOLS,
        )
        print("返回对象：")
        print(completion.choices[0].message.model_dump_json())
        print("\n")
        return completion

    def execute_tool(self, completion):
        print("正在执行工具函数...")
        function_name = completion.choices[0].message.tool_calls[0].function.name
        arguments_string = (
            completion.choices[0].message.tool_calls[0].function.arguments
        )
        arguments = json.loads(arguments_string)

        function = self.function_mapper[function_name]
        function_output = function() if arguments == {} else function(arguments)
        print(f"工具函数输出：{function_output}\n")
        return function_output


def main():
    assistant = Assistant()
    print("正在发起function calling...")
    completion = assistant.function_calling()
    assistant.execute_tool(completion)


if __name__ == "__main__":
    main()
